#coding:utf-8
#==============================================================================
#
"""
    Berkeley database implementation.
"""
#
#   Copyright (C) 2007 Michael A. Muller
#
#   This file is part of ODB.
#
#   ODB is free software: you can redistribute it and/or modify it under the 
#   terms of the GNU Lesser General Public License as published by the Free 
#   Software Foundation, either version 3 of the License, or (at your option) 
#   any later version.
#
#   ODB is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU Lesser General Public License for more details.
#
#   You should have received a copy of the GNU Lesser General Public License 
#   along with ODB.  If not, see <http://www.gnu.org/licenses/>.
#
#==============================================================================
 
import os, pickle, time, re
from threading import Lock, Thread
from StringIO import StringIO
from bsddb3.db import DBEnv as _DBEnv, DB as _DB
import bsddb3.db as _db
from odbi import DB, KeyDB, Env, Cursor, ReverseCursor

DO_NOT_INDEX = _db.DB_DONOTINDEX

_logRx = re.compile(r'log\.\d+')

class BDBDB(DB):
    """
        A database.
    """
 
    def __init__(self, env, name, type, allowDups):
        self.__env = env
        self._db = _DB(env)
        if allowDups:
  	    self._db.set_flags(_db.DB_DUP)
        self._db.open(name, dbtype = type, flags = _db.DB_CREATE | 
  		      _db.DB_THREAD | _db.DB_AUTO_COMMIT)
 
    def cursor(self, txn = None):
       """
           Returns a cursor on the database.
       """
       return BDBCursor(self._db, txn)
 
    def getTxn(self):
       """
 	   Convenience method - creates a transaction on the underlying
 	   environment and returns it.
       """
       return BDBTxn(self.__env.txn_begin())
 
class BDBCursor(Cursor):
    """
        A high-level cursor object.  Suitable for use in iteration.
 
        Cursors behave like python sequence indeces - think of them as pointing
        to the space between elements.  In this way, next() and prev() always
        have consistent meaning.  next() at the start of the table returns the
        first element, prev() at the end of the table returns the last element 
        and cursor.next() == cursor.prev().
    """
 
    def __init__(self, db, txn = None):
        """
            Creates a cursor from an underlying database object.  Not for user
            consumption, use @DB.cursor() to create a cursor.
        """
        self.__cursor = db.cursor(txn and txn._txn)
        self.__next = self.__prev = None
        if txn: txn._addCursor(self)
 
    def __convert(self, rawKeyVal):
        if rawKeyVal is None:
            return rawKeyVal
        return rawKeyVal[0], pickle.load(StringIO(rawKeyVal[1]))
 
    def setToFirst(self):
        """
            Sets the cursor to the first entry in the database.
        """
        self.__next = self.__convert(self.__cursor.first())
 
    def setToLast(self):
        self.__prev = self.__convert(self.__cursor.last())
 
    def setToKey(self, key, exact = False):
        """
  	    Sets the cursor to the record identified by the key.  The key is
  	    assumed to be a partial key unlss you specify otherwise.
  
  	    If an exact match is specified and the key is not matched, raises a
  	    *KeyError*
  
  	    parms:
  	        key::
  	            [string] the complete or partial key
  	        exact::
                    [boolean] if true, will match only an exact match.
                    Otherwise, matches the first key that is greater than or
                    equal to the key.
  
                    *Note,* this behavior only works on BTree's - on other
                    types of databases all matches will be exact.
        """
        if exact:
  	    self.__next = self.__convert(self.__cursor.set(key))
  	    if not self.__next:
  	        raise KeyError(key)
        else:
  	    self.__next = self.__convert(self.__cursor.set_range(key))
  
            # BDB cleverly sets the cursor to the first value in the table if
            # the key we pass to set_range() is greater than the last key in
            # the table.  We compensate for this greater intelligence by
            # forcing the cursor to the end of the table.
  	    if not self.__next:
  	        self.__prev = self.__convert(self.__cursor.last())
 
    def next(self):
        """
            Returns the next key/value pair referenced by the cursor and
            increments it.  Raises *StopIteration* if we're at the end of
            iteration.
        """
        if self.__next:
            temp = self.__next
            self.__next = None
        else:
            temp = self.__convert(self.__cursor.next())
            if temp is None:
                raise StopIteration()
        self.__prev = temp
        return temp
 
    def prev(self):
        """
            Returns the previous key/value pair referenced by the cursor and
            decrements it.  Raises *StopIteration* if we're at the end of
            iteration.
        """
        if self.__prev:
            temp = self.__prev
            self.__prev = None
        else:
            temp = self.__convert(self.__cursor.prev())
            if temp is None:
                raise StopIteration()
        self.__next = temp
        return temp
 
    def __iter__(self):
        """
            Allows the cursor to be iterated in a for loop:
 
            {{
               for key, val in cursor:
                  ...
            }}
        """
        return self
 
    def close(self):
        """
            Closes the cursor.
        """
        if self.__cursor:
            self.__cursor.close()
 
    def reverse(self):
        """
           Returns a @ReverseCursor allowing you to iterate over the list
           backwards.  Example:
 
           {{
               for key, val in cursor.reverse():
                   ...
           }}
        """
        return ReverseCursor(self)

class BDBKeyDB(KeyDB):
    """
        A data base that is accessible by a key.  All of the standard
        dictionary access methods are supported.
    """
    #类的专用方法可以做到obj = BDBKeyDB() obj["a"]
    def __getitem__(self, key):
        val = self._db.get(key)
        if val is None:
            raise KeyError(key)
        return pickle.load(StringIO(val))
        #定义了一个set方法将指定key的数据先序列化到StringIO里面去。相当于是在内存我们创建了一个内存文件先把数据set进去。然后找到了这个KEY就可以提取出来
        #======数据从内存中提取出来之后就马上实现序列化处理！

    def get(self, key, default = None, txn = None):
        val = self._db.get(key, txn = txn and txn._txn)
        if val is None:
            return default
        else:
            return pickle.load(StringIO(val))

    def has_key(self, key):
        return self._db.get(key) is not None

    def __delitem__(self, key):
        self.delete(None, key)

    def delete(self, txn, key):
        self._db.delete(key, txn and txn._txn)

    def __setitem__(self, key, val):
        self.set(None, key, val)

    def set(self, txn, key, val):
        buf = StringIO()
        pickle.dump(val, buf)
        flattened = buf.getvalue()
        self._db.put(key, flattened, txn and txn._txn)

class BDBBTreeDB(BDBDB, BDBKeyDB):

    def __init__(self, env, name, allowDups):
        BDBDB.__init__(self, env, name, _db.DB_BTREE, allowDups)

class BDBHashDB(BDBDB, BDBKeyDB):

    def __init__(self, env, name, allowDups):
        BDBDB.__init__(self, env, name, _db.DB_HASH, allowDups)

class BDBEnv(Env):
    """
       A database environment.  Groups together a set of databases.
    """

    def __init__(self, path, recover = False):

        # create the directory if it doesn't exist
        if not os.path.exists(path):
            os.mkdir(path)
        self.__env = _DBEnv()
        self.__env.open(path, 
                        _db.DB_INIT_LOG | _db.DB_INIT_LOCK |
                         _db.DB_INIT_MPOOL | _db.DB_INIT_TXN | _db.DB_CREATE |
                         _db.DB_THREAD | (recover and _db.DB_RECOVER_FATAL)
                        )
        self.__env.set_flags(_db.DB_AUTO_COMMIT, 1)

        # keep a cache of the database objects that we've created to improve
        # performance
        self.__dbs = {}

    def getBTreeDB(self, name, allowDups = False):
        if not self.__dbs.has_key(name):
            self.__dbs[name] = BDBBTreeDB(self.__env, name, allowDups)
        return self.__dbs[name]

    def getHashDB(self, name, allowDups = False):
        if not self.__dbs.has_key(name):
            self.__dbs[name] = BDBHashDB(self.__env, name, allowDups)
        return self.__dbs[name]

    def getTxn(self):
        """
            Returns a new transaction for the environment.  This is a raw BDB
            transaction.  See the @Txn class for documentation.
        """
        return Txn(self.__env.txn_begin())

    def getInactiveLogs(self):
        """
            Returns a list<string> of the filenames of the inactive logfiles
            for archiving purposes.
        """
        return self.__env.log_archive()

    def getDatabaseFiles(self):
        """
            Returns a list<string> of the filenames of all of the database
            files for archiving purposes.
        """
        return self.__env.log_archive(_db.DB_ARCH_DATA)

    def getAllLogs(self):
        """
            Returns a list<string> of the filenames of all of the logfiles for
            archiving purposes.
        """
        return self.__env.log_archive(_db.DB_ARCH_LOG)

    def checkDeadlocks(self):
        """
           Run one iteration of the deadlock detection algorithm.
        """
        return self.__env.lock_detect(_db.DB_LOCK_DEFAULT)

    class __DeadlockThread(Thread):
        def __init__(self, env, interval, logger):
            self.__env = env
            self.__interval = interval
            self.__logger = logger
            Thread.__init__(self)

        def __log(self, message):
            if self.__logger:
                self.__logger(message)

        def run(self):
            try:
                self.__log('starting deadlock thread')
                while True:
                    time.sleep(self.__interval)
                    self.__log('doing deadlock detection')
                    count = self.__env.checkDeadlocks()
                    if count:
                        self.__log('broke %d deadlocks' % count)
            except Exception, ex:
                self.__log('Deadlock thread terminated with an exception: %s' %
                            ex
                           )
             
    def startDeadlockThread(self, interval, logger = None):
        """
            Start a thread to check for deadlocks.

            parms:
                interval::
                    [float] time to sleep between deadlock checks (in seconds)
                logger::
                    [callable<string>] a callback to be used for logging
                    status info
        """
        thread = self.__DeadlockThread(self, interval, logger)
        thread.setDaemon(True)
        thread.start()
    
    def checkpoint(self):
        """
            Implements @odbi.Env.checkpoint().  Does nothing - checkpoint in 
            BDB happens automatically.
        """
        pass

class BDBTxn:
    """
        A database transaction.  Can be passed to all of the update functions.

        This isn't the real transaction class: the real transaction objects
        are raw bsddb3 transactions.  This class exists purely for
        documentation of that interface.
    """

    def __init__(self, txn):
        self._txn = txn
        self.__cursors = []

    def _addCursor(self, cursor):
        """
            Adds a new cursor to the transaction.  The cursor will be closed
            when the transaction is to prevent the database from getting into
            a state where it needs recovery.
        """
        self.__cursors.append(cursor)

    def __closeAllCursors(self):
        for cur in self.__cursors:
            cur.close()

    def abort(self):
        "Aborts the transaction, rolling back all changes."
        self.__closeAllCursors()
        self._txn.abort()

    def commit(self, flags = 0):
        """
            Commits the transaction.
        """
        self.__closeAllCursors()
        self._txn.commit(flags)

    def id(self):
        """
            Returns the unique transaction id associated with the transaction.
        """
        return self._txn.id()

    def prepare(self, gid):
        """
            Initiates the beginning of a two-phase commit.

            parms:
                gid::
                    [string] A global identifier, unique among all processes
                    involved in the commit.  Must by `DB_XIDDATASIZE` bytes
                    long.
        """
        self._txn.prepare(gid)

# vim:set sw=4:
